import pandas as pd
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error, mean_squared_error
from utils import load_json_data

# 加载数据
df = load_json_data("历年教学评价指标及影响因素数据.json")
print("数据加载完成，数据形状:", df.shape)
print("\n前5行数据:")
print(df.head())

X = df[['teacher_exp', 'training_count', 'has_phd', 'student_teacher_ratio', 'funding_per_student']]
y = df['teaching_score']

# 数据集划分
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
print(f"\n数据集划分完成:")
print(f"训练集形状: X_train={X_train.shape}, y_train={y_train.shape}")
print(f"测试集形状: X_test={X_test.shape}, y_test={y_test.shape}")

# 训练模型
lr_model = LinearRegression()
dt_model = DecisionTreeRegressor(random_state=42)

lr_model.fit(X_train, y_train)
dt_model.fit(X_train, y_train)
print("\n模型训练完成")

# 预测与评估
y_pred_lr = lr_model.predict(X_test)
y_pred_dt = dt_model.predict(X_test)

# 评估指标
metrics = {
    "linear_regression": {
        "MAE": round(mean_absolute_error(y_test, y_pred_lr), 3),
        "RMSE": round(np.sqrt(mean_squared_error(y_test, y_pred_lr)), 3)
    },
    "decision_tree": {
        "MAE": round(mean_absolute_error(y_test, y_pred_dt), 3),
        "RMSE": round(np.sqrt(mean_squared_error(y_test, y_pred_dt)), 3)
    }
}

print("\n模型评估指标:")
print("线性回归模型:")
print(f"  MAE: {metrics['linear_regression']['MAE']}")
print(f"  RMSE: {metrics['linear_regression']['RMSE']}")
print("决策树模型:")
print(f"  MAE: {metrics['decision_tree']['MAE']}")
print(f"  RMSE: {metrics['decision_tree']['RMSE']}")

# 特征重要性（仅决策树）
feature_importance = pd.DataFrame({
    "feature": X.columns,
    "importance": dt_model.feature_importances_
}).sort_values("importance", ascending=False).head(10)

print("\n决策树特征重要性:")
print(feature_importance)

# 未来3个学期预测（模拟数据）
future_semesters = ["2024-Spring", "2024-Fall", "2025-Spring"]
future_data = pd.DataFrame({
    "teacher_exp": [15, 18, 20],
    "training_count": [3, 4, 5],
    "has_phd": [1, 1, 0],
    "student_teacher_ratio": [20, 18, 22],
    "funding_per_student": [12000, 13000, 11000]
})
print("\n未来学期预测输入数据:")
print(future_data)

future_pred_lr = lr_model.predict(future_data).round(2)
future_pred_dt = dt_model.predict(future_data).round(2)

print("\n未来学期预测结果:")
for i, semester in enumerate(future_semesters):
    print(f"{semester}:")
    print(f"  线性回归预测: {future_pred_lr[i]}")
    print(f"  决策树预测: {future_pred_dt[i]}")

# 封装预测结果
teaching_pred_result = {
    "historical_data": df[['semester', 'teaching_score']].to_dict("records"),
    "future_pred": {
        "semesters": future_semesters,
        "lr_pred": future_pred_lr.tolist(),
        "dt_pred": future_pred_dt.tolist()
    },
    "feature_importance": feature_importance.to_dict("records"),
    "metrics": metrics
}

print("\n最终结果已封装到 teaching_pred_result 字典中")
print(teaching_pred_result)